{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ce7f288-2d83-47ee-ac51-741126d4e3bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from embedding_atlas.widget import EmbeddingAtlasWidget\n",
    "from embedding_atlas.projection import compute_text_projection\n",
    "import pandas as pd\n",
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69f21f68-c20e-4055-b073-5ab6c8f89913",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load a dataset\n",
    "ds = load_dataset(\"james-burton/wine_reviews\", split=\"validation\")\n",
    "df = pd.DataFrame(ds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a927e7d-8dda-4d6b-ae88-60cceb6f9678",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute text embedding and projection of the embedding using Sentence Transformers by default\n",
    "compute_text_projection(df, text=\"description\", x=\"projection_x\", y=\"projection_y\", neighbors=\"neighbors\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a82b7b13-2024-4ff9-b96d-8b21048ca2cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the dataset with the Embedding Atlas widget\n",
    "w = EmbeddingAtlasWidget(df, text=\"description\", x=\"projection_x\", y=\"projection_y\", neighbors=\"neighbors\")\n",
    "w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e42251b5-0e9f-4191-81fd-64522118f1cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the selection from the widget as a dataframe\n",
    "w.selection()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05056603",
   "metadata": {},
   "source": [
    "Embedding Atlas supports running text embeddings using all the models supported by [LiteLLM](https://docs.litellm.ai/docs/embedding/supported_embedding).\n",
    "\n",
    "To run the example below please make sure to first install [Ollama](https://ollama.com/download) and run the following command to download the [nomic-embed-text](https://ollama.com/library/nomic-embed-text) model:\n",
    "\n",
    "```bash\n",
    "ollama pull nomic-embed-text\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d2f6448-3ce9-4529-a81a-9caf3376be57",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute text embedding and projection of the embedding using locally running Ollama API\n",
    "compute_text_projection(\n",
    "    df,\n",
    "    text=\"description\",\n",
    "    x=\"projection_x\",\n",
    "    y=\"projection_y\",\n",
    "    neighbors=\"neighbors\",\n",
    "    text_projector=\"litellm\",\n",
    "    api_base_url=\"http://localhost:11434\",\n",
    "    model=\"ollama/nomic-embed-text\",\n",
    "    batch_size=512,\n",
    "    # Running batches synchronously to avoid overwhelming the local API\n",
    "    sync=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07b618a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the dataset with the Embedding Atlas widget using nomic-embed-text embeddings served by Ollama API\n",
    "EmbeddingAtlasWidget(df, text=\"description\", x=\"projection_x\", y=\"projection_y\", neighbors=\"neighbors\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aabe196c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Async processing of batches raises error in environments with an already running event loop, so we must patch the loop\n",
    "!uv pip install nest-asyncio -q\n",
    "import nest_asyncio\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6048064a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute text embedding and projection of the embedding using OpenAI API\n",
    "compute_text_projection(\n",
    "    df,\n",
    "    text=\"description\",\n",
    "    x=\"projection_x\",\n",
    "    y=\"projection_y\",\n",
    "    neighbors=\"neighbors\",\n",
    "    text_projector=\"litellm\",\n",
    "    # Your OpenAI API key. You can omit this and set the OPENAI_API_KEY environment variable instead.\n",
    "    api_key=\"sk-xxx\",\n",
    "    model=\"openai/text-embedding-3-small\",\n",
    "    # OpenAI's limit is 300K input token per request, so batch size should be chosen accordingly given the average `text` item length\n",
    "    batch_size=1024,\n",
    "    # Since calling a remote API is I/O-bound, we can benefit from async processing. This is the default behavior.\n",
    "    sync=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa498a52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display the dataset with the Embedding Atlas widget using text-embedding-3-small embeddings served by OpenAI API\n",
    "EmbeddingAtlasWidget(df, text=\"description\", x=\"projection_x\", y=\"projection_y\", neighbors=\"neighbors\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "embedding-atlas",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
